#include <assert.h>
#include <vector>
#include <iostream>
#include <cmath>
#include <fcntl.h>
#include <sys/mman.h>

#include <sys/time.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ipc.h>	// ftok
#include <unistd.h>		// fork
#include <string.h>		// memcpy
#include <sys/wait.h>	// waitpid
#include <stdbool.h>
#ifdef AFFINITY
#include <sched.h>
#endif
extern "C" {
//#include <cblas.h>
//#include  <clapack.h>
#include "f2c.h"
#include "blaswrap.h"
#include "clapack.h"
#include "cblas.h"
//#include "../../../libLC_unif/channel.h"
#include "channel.h"
}
//#include "./cycle.h"

#define M_EOF 'Y'
int core_mapping[16] = { 8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6, 7 };
int nWorkers;
int matrixSize;
int streamLength;
int blockSize;

// generic worker
complex *factorize_matrix(complex *originalMatrix, integer matrixSize, integer blockSize)
{

    complex *l = originalMatrix;
    int i, k, j; // indexes used in loops

    // Number of blocks
    int n = matrixSize / blockSize;

    complex alpha;
    complex beta;

    alpha.r = -1.;
    alpha.i = 0.;
    beta.r = 1.;
    beta.i = 0.;

    integer info;

    // Begin factorization
    for (k = 0; k < n; k++)
    {
        //Factorize the main diagonal block L[k][k]
        //factor(l[k][k])
        cpotf2_("Upper",
                &blockSize,
                //lkk
                (complex *) l + (k * (blockSize * matrixSize)) + ((k * blockSize)),
                &matrixSize,
                &info);

        //Compute all the blocks of the column k
        //ctrsm(l[*][k])
        for (i = k + 1; i < n; i++)
        {
            ctrsm_("Left",                    //Left
                   "Upper",                   //Upper
                   "Conjugate transpose",     //Conjugate transpose
                   "Non-unit",                //NoUnit
                   &blockSize,
                   &blockSize,
                   (complex *) &beta,
                   //lkk,
                   (complex *) l + (k * (blockSize * matrixSize)) + ((k * blockSize)),
                   &matrixSize,
                   (complex *) l + (i * (blockSize * matrixSize)) + ((k * blockSize)),
                   //lik,
                   &matrixSize);
        }

        //update all the other blocks of the main diagonal
        for (j = k + 1; j < n; j++)
        {
            cherk_("Upper",                          //Upper
                   "Conjugate Transpose",            //Conjugate Transpose
                   &blockSize,
                   &blockSize,
                   (real *) &alpha.r,
                   //ljk,
                   (complex *) l + (j * (blockSize * matrixSize)) + ((k * blockSize)),
                   &matrixSize,
                   (real *) &beta.r,
                   //ljj,
                   (complex *) l + (j * (blockSize * matrixSize)) + ((j * blockSize)),
                   &matrixSize);

            //update all the inner blocks
            for (i = j + 1; i < n; i++)
            {
                cgemm_("Conjugate transpose",         //No Transpose
                       "No transpose",                //Conjugate Transpose
                       &blockSize,
                       &blockSize,
                       &blockSize,
                       (complex *) &alpha,
                       //ljk,
                       (complex *) l + (j * (blockSize * matrixSize)) + ((k * blockSize)),
                       &matrixSize,
                       //lik,
                       (complex *) l + (i * (blockSize * matrixSize)) + ((k * blockSize)),
                       &matrixSize,
                       (complex *) &beta,
                       //lij,
                       (complex *) l + (i * (blockSize * matrixSize)) + ((j * blockSize)),
                       &matrixSize);

            }	// End of i loop

        }	// End of j loop

    }	// End of k loop

    return l;
}

/*
    -----------------   Emitter     ----------------------
*/
void emitter(int keys[], complex *matrix){
    //sleep(3);
	int i,j;

    struct timeval tmp_t;
    long start_tmill=0, end_tmill=0;
#ifdef AFFINITY
	cpu_set_t mask;
	/* CPU_ZERO initializes all the bits in the mask to zero. */
	CPU_ZERO(&mask);
	/* CPU_SET sets only the bit corresponding to cpu. */
	CPU_SET(0, &mask);
	/* sched_setaffinity returns 0 in success */
	if (sched_setaffinity(0, sizeof(mask), &mask) == -1)
	{
		printf("WARNING: Could not set CPU Affinity, continuing...\n");
	}
#endif
    printf("Emitter: i run on %d\n",sched_getcpu());
	Channel *out = (Channel *)malloc(sizeof(Channel)*nWorkers);

	// Prepare Channels
	for (i = 0; i < nWorkers; i++) {
		out[i] = LCattach(keys[i]);
	}

	Communicator kp=initCommunicator(core_mapping);

	for(j=0;j<streamLength/nWorkers;j++)
	{
		for(i=0; i<nWorkers; i++){
		    #ifdef DEBUG
		    printf("--- Emitter --- sent matrix to worker %d\n",i);
		    #endif
		    //we take the time for the sends
		    /*if(j>0)
            {
                gettimeofday(&tmp_t,NULL);
                start_tmill=tmp_t.tv_sec*1000+(tmp_t.tv_usec/1000);
            }*/
		    LCsend(kp,out[i],(char *)matrix);
		    //in this particular case, there is no need to wait (since the matrix is always the same)
		    /*if(j>0)
		    {
		        gettimeofday(&tmp_t,NULL);
                end_tmill+=tmp_t.tv_sec*1000+(tmp_t.tv_usec/1000)-start_tmill;
		    }*/
		}
	}
    //printf("Emitter: average time for each send %f\n",((double)end_tmill)/(streamLength-nWorkers));
	//printf("--- Emitter ---finalize,  detach channels and exit");
	finalizeCommunicator(kp);
	// Detach Channels
	for (i = 0; i < nWorkers; i++) {
		LCdetatch(out[i]);
	}
}

/*
    -----------------   Collector   ----------------------
*/
void collector(key_t keys[]){
	int i;
    struct timeval tmp_t;
    long start_tmill=0, end_tmill=0;
	#ifdef AFFINITY
	cpu_set_t mask;
	/* CPU_ZERO initializes all the bits in the mask to zero. */
	CPU_ZERO(&mask);
	/* CPU_SET sets only the bit corresponding to cpu. */
	CPU_SET(7, &mask); //we set it to the last CPU
	/* sched_setaffinity returns 0 in success */
	if (sched_setaffinity(0, sizeof(mask), &mask) == -1)
	{
		printf("WARNING: Could not set CPU Affinity, continuing...\n");
	}
    #endif

    printf("Collector: i run on %d\n",sched_getcpu());
    // Prepare Channels
    Channel *in = (Channel *)malloc(sizeof(Channel)*nWorkers);

	for (i = 0; i < nWorkers; i++) {
		in[i] = LCattach(keys[i+nWorkers]);
	}
	//there is no need of communicator here
    //Communicator kp=initCommunicator(core_mapping);


	char *result=NULL;
	// Start receiving from workers
    int j;
	for(j=0;j<streamLength/nWorkers;j++)
	{
		for(i=0; i<nWorkers; i++){
		    #ifdef DEBUG
		    printf("Collector[%u]: ricevo i dati da: %d\n",getpid(),i);
		    #endif
			result = LCreceive(in[i]);
			#ifdef DEBUG
			printf("Collector: dati ricevuti\n");
			#endif
		}
		if(j==1) //start to take time (we forget about the first two iteration)
        {
            gettimeofday(&tmp_t,NULL);
            start_tmill=tmp_t.tv_sec*1000+(tmp_t.tv_usec/1000);
        }
	}
    //compute time
    gettimeofday(&tmp_t,NULL);
    end_tmill=tmp_t.tv_sec*1000+(tmp_t.tv_usec/1000)-start_tmill;
    double st=((double)end_tmill)/(streamLength-2*nWorkers);
    printf("Collettore finito....\n");
    printf("Time elapsed (msec): %ld\n",end_tmill);
    printf("Service time(msec): %f\n",st);
    //print the matrix
    /*const char *factorizationFileName = "./choleskySequential.txt";
    complex *factorizedMatrix=((complex *)(result));
    FILE *fp = fopen(factorizationFileName, "w+");
    if (fp == NULL) {
        perror("Error opening output file");
        exit(-1);
    }

    fprintf(fp, "Block algorithm farm version on Intel: L result matrix \n");

    for (int i = 0; i < matrixSize; i++) {
        int j;
        fprintf(fp, "[ ");
        for (j = 0; j <= i ; j++) {
            fprintf(fp, "% 6.3f ", factorizedMatrix[i * matrixSize + j].r);
            fprintf(fp, "% 6.3fi ", factorizedMatrix[i * matrixSize + j].i);
        }
        for ( ; j < matrixSize; j++) {
            fprintf(fp, "% 6.3f ", 0.);
            fprintf(fp, "% 6.3fi ", 0.);
        }

        fprintf(fp, " ]\n");
    }

    fprintf(fp, "\n");

    fclose(fp);*/
	/*printf("Colector: finalizing....!\n");
    finalizeCommunicator(kp);
    printf("Collector: finalized!\n");*/
	// Detach Channels
	for (i = 0; i < nWorkers; i++) {
		LCdestroy(in[i]);
	}

    FILE *fp = fopen("resultKP.dat", "a");
    fprintf(fp,"%d %f\n",2*nWorkers,st);
    fclose(fp);
	exit(0);
}

/*
-----------------    Worker ---------------
- elements is the number of matrices that the worker has to
    compute
*/
int worker(int id, key_t inkey, key_t outkey, int elements) {

	int i;

//	printf("My id is %d. Channel in key: %d. Channel out key: %d Elements: %d\n", id, inkey, outkey, elements);

    #ifdef AFFINITY
        cpu_set_t mask;
        /* CPU_ZERO initializes all the bits in the mask to zero. */
        CPU_ZERO(&mask);
        /* CPU_SET sets only the bit corresponding to cpu. */
        CPU_SET(id+1, &mask);
        /* sched_setaffinity returns 0 in success */
        if (sched_setaffinity(0, sizeof(mask), &mask) == -1)
        {
            printf("WARNING: Could not set CPU Affinity, continuing...\n");
        }
        printf("Worker %d: i run on %d\n",id,sched_getcpu());
    #endif

    // Attach to the channel
	Channel in = LCattach(inkey);
	Channel out = LCattach(outkey);
    Communicator kp=initCommunicator(core_mapping);

    //printf("--- Worker %d --- i have to compute %d elements\n",id,elements);

    Delegation d=NULL;

    for(i=0;i<elements;i++)
    {
        #ifdef DEBUG
        printf("--- Worker %d --- receiving\n",id);
        #endif
        complex *matrix=(complex *)LCreceive(in);
        #ifdef DEBUG
        printf("--- Worker %d --- matrix received\n",id);
        #endif
        matrix = factorize_matrix(matrix,  matrixSize, blockSize);
        //wait for the previous send to be completed (otherwise the receive will overwrite it)
        if(d!=NULL)
            LCwait(d);
        //send to collector
        d=LCsend(kp,out,(char *)matrix);
        #ifdef DEBUG
        printf("--- Worker %d --- matrix factorized and sent\n",id);
        #endif
    }
    finalizeCommunicator(kp);
    LCdestroy(in);
	LCdetatch(out);

	//printf("Worker %d. Closing...\n", id);

	exit(0);

}

int main(int argc, char * argv[])
{

    complex* originalMatrix;

    if (argc < 6)
    {
        std::cerr << "usage: "
                  << argv[0]
                  << " matrixSize blockSize inputfile nWorker streamLength\n";
        return -1;
    }

    matrixSize = atoi(argv[1]);
    blockSize = atoi(argv[2]);
    char *infile = argv[3];
    nWorkers=atoi(argv[4]);
    streamLength=atoi(argv[5]);

    /* Una restrizione che imponiamo e' che' il numero di elementi dello stream
        sia divisible per il numero di worker*/
    if(streamLength%nWorkers!=0)
    {
        fprintf(stderr, "Error: the stream length must be divisible for the number of workers\n");
        exit(-1);
    }
    if(nWorkers>6)
    {
        fprintf(stderr, "Error: max number of workers is 6\n");
        exit(-1);
    }

    Channel tmpchan;
    key_t *keys = (key_t *)(malloc(sizeof(key_t) * nWorkers * 2 ));
    /*
        Message: it will contain an initial character that says if this is the last element
        of the stream (sent by the emitter to a particular worker or by a worker to the collector).

        So it will have dimension 1+matrixSize * matrixSize * sizeof(complex)
    */

    /*
     * I create the channels here so they exists when workers and emitter starts
     * However they are immediately detached so they are not retained after forks
     */
     int i;
    for (i = 0; i < nWorkers; i++)
    {
        keys[i] = ftok(argv[0],i);
        keys[i+nWorkers] = ftok(argv[0],i+nWorkers);
        if(keys[i]==-1)
        {
            printf("Error creating a new key with parameters %s %d\n",argv[0],i);
            exit(-1);
        }

        tmpchan =  LCcreate_as( keys[i], (matrixSize * matrixSize * sizeof(complex)) , 1, 2 );
        LCdetatch(tmpchan);
        tmpchan =  LCcreate_as( keys[i+nWorkers], (matrixSize * matrixSize * sizeof(complex)), 1, 2 );
        LCdetatch(tmpchan);
    }

    /*
        Creation of the processes
    */


     int *pids=(int *)malloc(sizeof(int)*(nWorkers+1));
    for (i = 0; i < nWorkers; i++)
    {
        pid_t pid = fork();
        if(pid==0)  // Child
        {
            worker(i,keys[i], keys[i+nWorkers],streamLength/nWorkers);
        }
        else if (pid < 0)   // failed to fork
        {
            printf("Failed to fork\n");
            exit(1);
            // Throw exception
        }
        else
        {
            pids[i] = pid;
        }
    }
    // Yet another fork, for the collector
    pid_t pid;
    {
        pid = fork();
        if(pid==0)  // Child
        {
            collector(keys);
        }
        else if (pid < 0)   // failed to fork
        {
            printf("Failed to fork\n");
            exit(1);
            // Throw exception
        }
        else
        {
            pids[i] = pid;
        }
    }

    FILE *fp = fopen(infile, "r");
    if (!fp)
    {
        std::cerr << "Error opening input file\n";
        return -1;
    }

    // Memory allocation without HUGE_PAGES
    originalMatrix = (complex *) malloc(matrixSize * matrixSize * sizeof(complex));

    // Reads matrix A from input file and initializes matrix L
    int n;
    fscanf(fp, "%d", &n);
    assert(n == matrixSize);
    for (int i = 0; i < n; i++)
    {
        for (int j = 0; j < n; j++)
        {
            fscanf(fp, "%f", &(originalMatrix[(i * matrixSize) + j].r));
            fscanf(fp, "%f", &(originalMatrix[(i * matrixSize) + j].i));
        }
    }

    fclose(fp);

    //act as emitter
    emitter(keys,originalMatrix);

    // Wait for the workers to quit
    int childExitStatus;
	for (i = 0; i < nWorkers; i++) {
        printf("Waitpid %d \n",i);
		waitpid( pids[i] , &childExitStatus, 0);
		printf("Waitpid %d conclusa\n",i);
	}
	// Wait for the collector to quit
	waitpid(pid, &childExitStatus, 0);

	return 0;
}
